#!/bin/bash
# AWS启动脚本 - 在AWS节点上启动Turritopsis协议

set -e

# 配置参数
KEY_NAME="turritopsis-key"
PROJECT_DIR="/home/ec2-user/turritopsis"
EPOCHS=30
R=5
N=9
F=2
CURVE="secp256k1"
VERIFY_BATCH=true
VERIFY_COUNT=128
BATCH_SIZE=1000

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_debug() {
    echo -e "${BLUE}[DEBUG]${NC} $1"
}

# 检查hosts文件
check_hosts_file() {
    if [ ! -f "aws_hosts.txt" ]; then
        log_error "aws_hosts.txt not found. Please run aws_deploy.sh first."
        exit 1
    fi
    
    log_info "Found $(wc -l < aws_hosts.txt) nodes in aws_hosts.txt"
}

# 生成节点配置
generate_node_configs() {
    log_info "Generating node configurations..."
    
    # 创建配置目录
    mkdir -p aws_configs
    
    node_id=0
    while read -r node_name public_ip private_ip; do
        if [[ "$node_name" =~ ^node- ]]; then
            # 生成节点配置
            cat > "aws_configs/node_${node_id}_config.yaml" << EOF
# Turritopsis Node $node_id Configuration
node_id: $node_id
public_ip: $public_ip
private_ip: $private_ip
epochs: $EPOCHS
r: $R
n: $N
f: $F
curve: $CURVE
verify_batch: $VERIFY_BATCH
verify_count: $VERIFY_COUNT
batch_size: $BATCH_SIZE
router_kind: tcp
router_addr: $private_ip:19010
use_router: true
EOF
            ((node_id++))
        fi
    done < aws_hosts.txt
    
    log_info "Generated $node_id node configurations"
}

# 启动ACS stub服务器
start_acs_servers() {
    log_info "Starting ACS stub servers on all nodes..."
    
    node_id=0
    while read -r node_name public_ip private_ip; do
        if [[ "$node_name" =~ ^node- ]]; then
            log_info "Starting ACS server on $node_name ($public_ip)..."
            
            # 创建启动脚本
            cat > "start_acs_${node_id}.sh" << EOF
#!/bin/bash
cd $PROJECT_DIR
nohup ./acs_stub_server > acs_server.log 2>&1 &
echo \$! > acs_server.pid
echo "ACS server started on node $node_id"
EOF
            
            # 复制并执行
            scp -i ~/.ssh/$KEY_NAME.pem -o StrictHostKeyChecking=no "start_acs_${node_id}.sh" ec2-user@$public_ip:/tmp/
            ssh -i ~/.ssh/$KEY_NAME.pem -o StrictHostKeyChecking=no ec2-user@$public_ip \
                "chmod +x /tmp/start_acs_${node_id}.sh && /tmp/start_acs_${node_id}.sh"
            
            ((node_id++))
        fi
    done < aws_hosts.txt
    
    log_info "ACS servers started on all nodes"
}

# 启动TCP路由器
start_tcp_routers() {
    log_info "Starting TCP routers on all nodes..."
    
    node_id=0
    while read -r node_name public_ip private_ip; do
        if [[ "$node_name" =~ ^node- ]]; then
            log_info "Starting TCP router on $node_name ($public_ip)..."
            
            # 创建启动脚本
            cat > "start_router_${node_id}.sh" << EOF
#!/bin/bash
cd $PROJECT_DIR
nohup ./turritopsis_full --router tcp --router-addr $private_ip:19010 --router-server true > router.log 2>&1 &
echo \$! > router.pid
echo "TCP router started on node $node_id"
EOF
            
            # 复制并执行
            scp -i ~/.ssh/$KEY_NAME.pem -o StrictHostKeyChecking=no "start_router_${node_id}.sh" ec2-user@$public_ip:/tmp/
            ssh -i ~/.ssh/$KEY_NAME.pem -o StrictHostKeyChecking=no ec2-user@$public_ip \
                "chmod +x /tmp/start_router_${node_id}.sh && /tmp/start_router_${node_id}.sh"
            
            ((node_id++))
        fi
    done < aws_hosts.txt
    
    log_info "TCP routers started on all nodes"
}

# 启动Turritopsis节点
start_turritopsis_nodes() {
    log_info "Starting Turritopsis nodes..."
    
    # 等待路由器启动
    log_info "Waiting for routers to start..."
    sleep 10
    
    node_id=0
    while read -r node_name public_ip private_ip; do
        if [[ "$node_name" =~ ^node- ]]; then
            log_info "Starting Turritopsis on $node_name ($public_ip)..."
            
            # 创建启动脚本
            cat > "start_turritopsis_${node_id}.sh" << EOF
#!/bin/bash
cd $PROJECT_DIR

# 设置环境变量
export ACS_BATCH_SIZE=$BATCH_SIZE

# 启动Turritopsis
nohup ./turritopsis_full \\
    --epochs $EPOCHS \\
    --r $R \\
    --n $N \\
    --f $F \\
    --pid $node_id \\
    --curve $CURVE \\
    --verify-batch $VERIFY_BATCH \\
    --verify-count $VERIFY_COUNT \\
    --batch-size $BATCH_SIZE \\
    --router tcp \\
    --router-addr $private_ip:19010 \\
    --use-router true \\
    --acs-addr $private_ip:19001 \\
    > turritopsis.log 2>&1 &

echo \$! > turritopsis.pid
echo "Turritopsis started on node $node_id (PID: \$(cat turritopsis.pid))"
EOF
            
            # 复制并执行
            scp -i ~/.ssh/$KEY_NAME.pem -o StrictHostKeyChecking=no "start_turritopsis_${node_id}.sh" ec2-user@$public_ip:/tmp/
            ssh -i ~/.ssh/$KEY_NAME.pem -o StrictHostKeyChecking=no ec2-user@$public_ip \
                "chmod +x /tmp/start_turritopsis_${node_id}.sh && /tmp/start_turritopsis_${node_id}.sh"
            
            ((node_id++))
        fi
    done < aws_hosts.txt
    
    log_info "Turritopsis nodes started on all nodes"
}

# 监控节点状态
monitor_nodes() {
    log_info "Monitoring node status..."
    
    while true; do
        echo "=== Node Status $(date) ==="
        
        node_id=0
        while read -r node_name public_ip private_ip; do
            if [[ "$node_name" =~ ^node- ]]; then
                # 检查进程状态
                status=$(ssh -i ~/.ssh/$KEY_NAME.pem -o StrictHostKeyChecking=no -o ConnectTimeout=5 ec2-user@$public_ip \
                    "if [ -f $PROJECT_DIR/turritopsis.pid ]; then ps -p \$(cat $PROJECT_DIR/turritopsis.pid) > /dev/null 2>&1 && echo 'RUNNING' || echo 'STOPPED'; else echo 'NO_PID'; fi" 2>/dev/null || echo "ERROR")
                
                echo "Node $node_id ($public_ip): $status"
                ((node_id++))
            fi
        done < aws_hosts.txt
        
        echo ""
        sleep 30
    done
}

# 停止所有节点
stop_all_nodes() {
    log_info "Stopping all nodes..."
    
    node_id=0
    while read -r node_name public_ip private_ip; do
        if [[ "$node_name" =~ ^node- ]]; then
            log_info "Stopping $node_name ($public_ip)..."
            
            # 停止Turritopsis
            ssh -i ~/.ssh/$KEY_NAME.pem -o StrictHostKeyChecking=no ec2-user@$public_ip \
                "cd $PROJECT_DIR && if [ -f turritopsis.pid ]; then kill \$(cat turritopsis.pid) 2>/dev/null; rm -f turritopsis.pid; fi"
            
            # 停止路由器
            ssh -i ~/.ssh/$KEY_NAME.pem -o StrictHostKeyChecking=no ec2-user@$public_ip \
                "cd $PROJECT_DIR && if [ -f router.pid ]; then kill \$(cat router.pid) 2>/dev/null; rm -f router.pid; fi"
            
            # 停止ACS服务器
            ssh -i ~/.ssh/$KEY_NAME.pem -o StrictHostKeyChecking=no ec2-user@$public_ip \
                "cd $PROJECT_DIR && if [ -f acs_server.pid ]; then kill \$(cat acs_server.pid) 2>/dev/null; rm -f acs_server.pid; fi"
            
            ((node_id++))
        fi
    done < aws_hosts.txt
    
    log_info "All nodes stopped"
}

# 主函数
main() {
    log_info "Starting Turritopsis on AWS nodes..."
    
    check_hosts_file
    generate_node_configs
    start_acs_servers
    start_tcp_routers
    start_turritopsis_nodes
    
    log_info "Turritopsis started on all nodes!"
    log_info "Use 'monitor' command to monitor status"
    log_info "Use 'stop' command to stop all nodes"
}

# 脚本参数处理
case "${1:-}" in
    "start")
        main
        ;;
    "monitor")
        monitor_nodes
        ;;
    "stop")
        stop_all_nodes
        ;;
    "status")
        log_info "Checking node status..."
        node_id=0
        while read -r node_name public_ip private_ip; do
            if [[ "$node_name" =~ ^node- ]]; then
                status=$(ssh -i ~/.ssh/$KEY_NAME.pem -o StrictHostKeyChecking=no -o ConnectTimeout=5 ec2-user@$public_ip \
                    "if [ -f $PROJECT_DIR/turritopsis.pid ]; then ps -p \$(cat $PROJECT_DIR/turritopsis.pid) > /dev/null 2>&1 && echo 'RUNNING' || echo 'STOPPED'; else echo 'NO_PID'; fi" 2>/dev/null || echo "ERROR")
                echo "Node $node_id ($public_ip): $status"
                ((node_id++))
            fi
        done < aws_hosts.txt
        ;;
    *)
        echo "Usage: $0 {start|monitor|stop|status}"
        echo "  start   - Start Turritopsis on all nodes"
        echo "  monitor - Monitor node status continuously"
        echo "  stop    - Stop all nodes"
        echo "  status  - Check current status"
        exit 1
        ;;
esac